#!/usr/bin/env python3
# set coding: utf-8
"""
@Project : litewaf
@File :    init.py
@Time :    2023/3/9 21:32
@Author :  richard zhu
@Email :   gaotao.zhugt@dtzhejiang.com
@purpose : 加载默认配置
"""
import random

from django.core.management.base import CommandError, BaseCommand
from api.models import GlobalSetting, RuleGroup, WafRule, FlowRule, WafUser
from django.contrib.auth.models import User
import string


def gen_passwd(len=12):
    letters = string.digits + string.ascii_letters
    passwd = ""
    for i in range(len):
        passwd += random.choice(letters)
    return passwd

def web_block(block_k, block_v, rule_strings):
    try:
        block_name = block_k + "_" + block_v
        rg = RuleGroup.objects.get_or_create(name="%s默认防护规则" % block_name)
        if rg[1]:
            print("%s规则组创建成功" % block_name)
        else:
            print("%s规则组已存在" % block_name)


        for ii,i in enumerate(rule_strings):
            rule_matchs = {
                "match_args": [{"key": block_k, "value": block_v}],
                "args_prepocess": ["none"],
                "match_operator": "str_contain",
                "match_value": i
            }
            rule = WafRule.objects.get_or_create(rule_name="%s_%s" % (block_name, ii) ,
                                           defaults={
                                               "rule_action":"block",
                                               "rule_matchs":[rule_matchs]
                                            },
                                            )
            if not rule[1]:
                print("rule_name[%s_%s]已存在" % (block_name, ii))
            rg[0].waf_rule.add(rule[0])
    except Exception as e:
        print(e)

def flow_limit():
    rg = RuleGroup.objects.get_or_create(name="限流默认防护规则")
    if rg[1]:
        print("限流规则组创建成功")
    else:
        print("限流规则组已存在")

    f1 = FlowRule.objects.get_or_create(rule_name='参数限流1-优先级较高-对指定域名整体流量进行保护',
                                       rule_type='flow_limit_burst_rule',
                                       action_value={"rate": 2, "burst": 1},
                                       rule_action='burst_limit_redis',
                                       rule_matchs=[{"match_args": [{"key": "http_args", "value": "host"}], "args_prepocess": ["none"], "match_operator": "str_contain", "match_value": "ip.gaotao.club"}],
                                       flow_factor_keys={"key": "http_args", "value": "src_ip"},  # 针对字段进行聚合统计，可以是源地址，也可以是域名
                                       order=5)  # 优先级低

    f2 = FlowRule.objects.get_or_create(rule_name='参数限流2-优先级高-限制特定IP的并发流量',
                                       rule_type='flow_limit_burst_rule',
                                       action_value={"rate": 20, "burst": 10},
                                       rule_action='burst_limit_redis',
                                       rule_matchs=[{"match_args": [{"key": "http_args", "value": "host"}], "args_prepocess": ["none"], "match_operator": "str_contain", "match_value": "ip.gaotao.club"}, {"match_args": [{"key": "http_args", "value": "src_ip"}], "args_prepocess": ["none"], "match_operator": "str_contain", "match_value": "192.168.218.1"}],
                                       flow_factor_keys={},
                                       order=2) # 优先级高，对特定参数条件做限制，优先级一定要比聚合的高，先匹配这条，在匹配聚合的

    f3 = FlowRule.objects.get_or_create(rule_name='流量保护1-优先级低-对请求数量进行单位时间内的限制',
                                        rule_type='flow_protection_rule',
                                        action_value={"limit": 20, "time": 60},
                                        rule_action='flow_protection_local',
                                        rule_matchs=[{"match_args": [{"key": "http_args", "value": "host"}], "args_prepocess": ["none"], "match_operator": "str_contain", "match_value": "ip.gaotao.club"}],
                                        flow_factor_keys={},
                                        order=10)
    if not f1[1]:
        print("f1已存在" )
    else:
        rg[0].flow_rule.add(f1[0])

    if not f2[1]:
        print("f2已存在" )
    else:
        rg[0].flow_rule.add(f2[0])

    if not f3[1]:
        print("f3已存在" )
    else:
        rg[0].flow_rule.add(f3[0])

class Command(BaseCommand):
    help = '初始化litewaf'  # command功能作用简介

    def handle(self, *args, **kwargs):  # 主处理程序
        # 新增用户

        try:
            gs = GlobalSetting.objects.get_or_create(
                    name="global_setting",
                    defaults={
                        "default_page":{"code": "404", "html": "Not in default ServerName"},
                        "redis":{"host": "127.0.0.1", "port": 6379, "password": ""},
                        "sys_abnormal_handle_data":{"bypass_check": True, "same_name_args_check": True, "truncated_agrs_check": True,
                                                  "client_body_size_check": True, "code": "401", "html": "litewaf request abnormal handle"},
                        "waf_module_enable":["waf_rule_check", "domain_check", "flow_protection_check", "flow_limit_burst_check"],
                        "auto_update_period":30,
                        "flow_protection":{"default_limit": 1000, "code": "401", "html": "<html><head><meta http-equiv=\"Content-Type\" content=\"text/html; charset=utf-8\" /><meta http-equiv=\"Content-Language\" content=\"zh-cn\" /><title>网站防火墙</title></head><body><h1 align=\"center\"> 网站waf防火墙已拦截 </body></html>"},
                        "block_page":{"code": "403", "html": "<html><head><meta http-equiv=\"Content-Type\" content=\"text/html; charset=utf-8\" /><meta http-equiv=\"Content-Language\" content=\"zh-cn\" /><title>网站防火墙</title></head><body><h1 align=\"center\"> litewaf防火墙已拦截 </body></html>"}
                    }
                )
            if gs[1]:
                print("globalSetting初始化成功")
            else:
                print("globalSetting 已存在")

            web_block(block_k="http_args", block_v="user_agent",
                      rule_strings=["HTTrack|harvest|audit|dirbuster|pangolin|nmap|sqln|-scan|hydra|Parser|libwww|BBBike|sqlmap|w3af|owasp|Nikto|fimap|havij|PycURL|zmeu|BabyKrokodil|netsparker|httperf|bench"]
                     )

            rule_strings = ['\.\./',
                            '\:\$',
                            '\$\{',
                            'select.+(from|limit)',
                            '(?:(union(.*?)select))',
                            'having|rongjitest',
                            'sleep\((\s*)(\d*)(\s*)\)',
                            'benchmark\((.*)\,(.*)\)'      
                            'base64_decode\(',
                            '(?:from\W+information_schema\W)',
                            '(?:(?:current_)user|database|schema|connection_id)\s*\(',
                            '(?:etc\/\W*passwd)',
                            'into(\s+)+(?:dump|out)file\s*',
                            'group\s+by.+\(',
                            'xwork.MethodAccessor',
                            '(?:define|eval|file_get_contents|include|require|require_once|shell_exec|phpinfo|system|passthru|preg_\w+|execute|echo|print|print_r|var_dump|(fp)open|alert|showmodaldialog)\(',
                            'xwork\.MethodAccessor',
                            '(gopher|doc|php|glob|file|phar|zlib|ftp|ldap|dict|ogg|data)\:\/',
                            'java\.lang',
                            '\$_(GET|post|cookie|files|session|env|phplib|GLOBALS|SERVER)\[',
                            '\<(iframe|script|body|img|layer|div|meta|style|base|object|input)',
                            '(onmouseover|onerror|onload)\='
                            ]

            web_block(block_k="http_args",block_v="query_string",rule_strings=rule_strings)
            web_block(block_k="http_args",block_v="raw_body",rule_strings=rule_strings)

            url_list = [
                '\.(htaccess|bash_history)',
                '\.(bak|inc|old|mdb|sql|backup|java|class|tgz|gz|tar|zip)$',
                '(phpmyadmin|jmx-console|admin-console|jmxinvokerservlet)',
                'java\.lang',
                '\.svn\/',
                '/(attachments|upimg|images|css|uploadfiles|html|uploads|templets|static|template|data|inc|forumdata|upload|includes|cache|avatar)/(\\w+).(php|jsp)'
            ]


            web_block(block_k="http_args",
                      block_v="path",
                      rule_strings=url_list)

            cookie_list = [
                '\.\./',
                '\:\$',
                '\$\{',
                'select.+(from|limit)',
                '(?:(union(.*?)select))',
                'having|rongjitest',
                'sleep\((\s*)(\d*)(\s*)\)',
                'benchmark\((.*)\,(.*)\)',
                'base64_decode\(',
                '(?:from\W+information_schema\W)',
                '(?:(?:current_)user|database|schema|connection_id)\s*\(',
                '(?:etc\/\W*passwd)',
                'into(\s+)+(?:dump|out)file\s*',
                'group\s+by.+\(',
                'xwork.MethodAccessor',
                '(?:define|eval|file_get_contents|include|require|require_once|shell_exec|phpinfo|system|passthru|preg_\w+|execute|echo|print|print_r|var_dump|(fp)open|alert|showmodaldialog)\(',
                'xwork\.MethodAccessor',
                '(gopher|doc|php|glob|file|phar|zlib|ftp|ldap|dict|ogg|data)\:\/',
                'java\.lang',
                '\$_(GET|post|cookie|files|session|env|phplib|GLOBALS|SERVER)\['
            ]

            web_block(block_k="header_args",
                      block_v="cookie",
                      rule_strings=cookie_list)


            post_kv_check = {
                'password':'123456'
            }
            for k,v in post_kv_check.items():
                web_block(block_k="post_args",
                          block_v=k,
                          rule_strings=[v])


            flow_limit()
        except Exception as e:
            print(str(e))

        try:
            passwd = gen_passwd()
            User.objects.create_superuser(username='litewaf', password=passwd)
        except Exception as e:
            print("创建管理员账号失败",e)
        else:
            print("管理员用户litewaf创建成功，密码为[%s]" % passwd)

        try:
            passwd = gen_passwd()
            WafUser.objects.create(user_name='client', passwd=passwd)
        except Exception as e:
            print("创建litewaf客户端账号失败",e)
        else:
            print("客户端用户litewaf创建成功，密码为[%s]" % passwd)






